{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "62259398-4cb0-4ae3-9a84-87ed649b0a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shap\n",
    "from sklearn.tree import DecisionTreeRegressor, plot_tree, _tree\n",
    "from matplotlib import pyplot as plt\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import  train_test_split\n",
    "import pickle\n",
    "import xgboost\n",
    "import numpy as np\n",
    "from sklearn import preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "56e87bef-c1f4-42f4-b25f-ec83b4b49b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_df():\n",
    "    df = pd.read_csv(\"./data/titanic.csv\")\n",
    "    df = df[[\"Pclass\", \"Sex\", \"Age\", \"SibSp\", \"Parch\",\"Fare\", \"Cabin\", \"Embarked\", \"Survived\"]]\n",
    "    df[\"Cabin\"] = [x[0] if type(x) == str else \"No Data\" for x in df[\"Cabin\"].values]\n",
    "    df = df.dropna().reset_index(drop=True)\n",
    "\n",
    "    for column in [\"Sex\", \"Cabin\", \"Embarked\"]:\n",
    "        le = preprocessing.LabelEncoder()\n",
    "        le.fit(df[column].values)\n",
    "        df[column] = le.transform(df[column])\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "27867c91-85d2-4952-a774-96978e848d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_X_y(df):\n",
    "    X = df.drop(\"Survived\", axis=1)\n",
    "    y = df[\"Survived\"]\n",
    "    return X, y\n",
    "\n",
    "def get_model(X, y):\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "    model_xgb = xgboost.XGBClassifier(n_estimators=100, max_depth=2).fit(X_train, y_train)\n",
    "    return model_xgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "0f1c42f3-80e4-4824-abbf-b04b3047cdb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_decition_tree(X, y, depth):\n",
    "    clf = DecisionTreeRegressor(max_depth=depth)\n",
    "    clf = clf.fit(X, y)\n",
    "\n",
    "    return clf\n",
    "\n",
    "def get_shap_values_for_tree(model, X, y):\n",
    "    explainer = shap.TreeExplainer(model)\n",
    "    shap_values = explainer.shap_values(X, y=y)\n",
    "    return shap_values\n",
    "\n",
    "def get_y_proba_for_tree(model, X):\n",
    "    return model.predict_proba(X)[:,1]\n",
    "\n",
    "def get_cleaned_data_for_shap(df):\n",
    "    clean_df = df.dropna()\n",
    "    X = clean_df.drop(\"Survived\", axis=1)\n",
    "    y = clean_df[\"Survived\"]\n",
    "    return X, y\n",
    "\n",
    "def get_feature_index(feature_name, X):\n",
    "    return list(X.columns).index(feature_name)\n",
    "\n",
    "def get_feature_shap_vals(shap_values, feature_index):\n",
    "    return shap_values[:,feature_index]\n",
    "\n",
    "def plot_fig_for_tree(clf, feature_names):\n",
    "    fig = plt.figure(figsize=(15,5))\n",
    "    _ = plot_tree(clf, filled=True, feature_names=feature_names)\n",
    "    \n",
    "\n",
    "def get_decition_tree_for_feature(model, df, feature_name, depth, should_remove_original_feature=False):\n",
    "    X, y = get_cleaned_data_for_shap(df)\n",
    "    shap_values = get_shap_values_for_tree(model, X, y)\n",
    "\n",
    "    feature_index = get_feature_index(feature_name, X)\n",
    "    feature_shap_vals = get_feature_shap_vals(shap_values, feature_index)\n",
    "    if should_remove_original_feature:\n",
    "        X = X.drop(feature_name, axis=1)\n",
    "\n",
    "    clf = fit_decition_tree(X, feature_shap_vals, depth)\n",
    "\n",
    "    plot_fig_for_tree(clf, X.columns)\n",
    "\n",
    "def get_decition_tree_text_for_feature(model, df, feature_name, depth, should_remove_original_feature=False):\n",
    "    X, y = get_cleaned_data_for_shap(df)\n",
    "    shap_values = get_shap_values_for_tree(model, X, y)\n",
    "    \n",
    "    feature_index = get_feature_index(feature_name, X)\n",
    "    feature_shap_vals = get_feature_shap_vals(shap_values, feature_index)   \n",
    "    \n",
    "    if should_remove_original_feature:\n",
    "        X = X.drop(feature_name, axis=1)\n",
    "\n",
    "    clf = fit_decition_tree(X, feature_shap_vals, depth)\n",
    "\n",
    "    return clf, X\n",
    "\n",
    "def get_decition_tree_model(model, df, depth):\n",
    "    X, y = get_cleaned_data_for_shap(df)\n",
    "    y_proba = get_y_proba_for_tree(model, X)\n",
    "    \n",
    "    clf = fit_decition_tree(X, y_proba, depth)\n",
    "\n",
    "    return clf, X\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "8e3808d2-9497-4b35-9ee0-df3c05206dad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_rules(tree, feature_names, index, res_name=\"Важность\"):\n",
    "    tree_ = tree.tree_\n",
    "    feature_name = [feature_names[i] for i in tree_.feature]\n",
    "\n",
    "    paths = []\n",
    "    path = []\n",
    "    \n",
    "    def recurse(node, path, paths):\n",
    "        \n",
    "        if tree_.feature[node] != _tree.TREE_UNDEFINED:\n",
    "            name = feature_name[node]\n",
    "            threshold = tree_.threshold[node]\n",
    "            p1, p2 = list(path), list(path)\n",
    "            p1 += [f\"({name} <= {np.round(threshold, 3)})\"]\n",
    "            recurse(tree_.children_left[node], p1, paths)\n",
    "            p2 += [f\"({name} > {np.round(threshold, 3)})\"]\n",
    "            recurse(tree_.children_right[node], p2, paths)\n",
    "        else:\n",
    "            path += [(tree_.value[node], tree_.n_node_samples[node])]\n",
    "            paths += [path]\n",
    "            \n",
    "    recurse(0, path, paths)\n",
    "\n",
    "    samples_count = [p[-1][1] for p in paths]\n",
    "    ii = list(np.argsort(samples_count))\n",
    "    paths = [paths[i] for i in reversed(ii)]\n",
    "    \n",
    "    rules = []\n",
    "    for path in paths:\n",
    "        rule = \"Если \"\n",
    "        \n",
    "        for p in path[:-1]:\n",
    "            if rule != \"Если \":\n",
    "                rule += \" И \"\n",
    "            rule += f\"{p}\"\n",
    "        val = int(np.round(path[-1][0][0][0] * 100,0))\n",
    "        rule = f\"{res_name}={val} ({rule})\"\n",
    "        \n",
    "        \n",
    "        if index:\n",
    "            rule += f\" Тогда ({res_name}_{index}={np.round(path[-1][0][0][0],3)})\"\n",
    "        rules += [rule]\n",
    "        \n",
    "    return rules\n",
    "\n",
    "def get_rules_for_features(model_xgb, df, features, depth):\n",
    "    rules = []\n",
    "\n",
    "    for i, f in enumerate(features):\n",
    "        clf, X = get_decition_tree_text_for_feature(model_xgb, df, f, depth)\n",
    "        rules += get_rules(clf, X.columns, i)\n",
    "\n",
    "    return rules\n",
    "\n",
    "def get_rules_for_y_proba(model_xgb, df, depth):\n",
    "    clf, X = get_decition_tree_model(model_xgb, df, depth)\n",
    "    rules = get_rules(clf, X.columns, None, \"Proba\")\n",
    "\n",
    "    return rules\n",
    "\n",
    "def format_rules_as_str(rules):\n",
    "    return \"\\n\".join(rules)\n",
    "\n",
    "def fix_comparison(rules_str):\n",
    "    replace_rules = {\n",
    "        \"Sex > 0.5\": \"Sex = M\",\n",
    "        \"Sex <= 0.5\": \"Sex = F\",\n",
    "        \"Pclass > 2.5\": \"Pclass = High\",\n",
    "        \"Pclass <= 2.5\": \"Pclass = Low\",\n",
    "        \"Вероятность_0\": \"Вероятность\"\n",
    "    }\n",
    "    for from_, to in replace_rules.items():\n",
    "        rules_str = rules_str.replace(from_, to)\n",
    "\n",
    "    return rules_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "264e6a93-1631-4222-8d30-2ff313a6a2ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = get_df()\n",
    "X, y = get_X_y(df)\n",
    "model = get_model(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "5463568d-d2de-4af7-88bd-dd8a1ff219e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Proba=12 (Если (Sex = M) И (Cabin > 6.0) И (Age > 3.5))\n",
      "Proba=41 (Если (Sex = M) И (Cabin <= 6.0) И (Age > 14.0))\n",
      "Proba=88 (Если (Sex = F) И (Pclass = Low) И (Fare <= 51.74))\n",
      "Proba=56 (Если (Sex = F) И (Pclass = High) И (Fare <= 20.8))\n",
      "Proba=97 (Если (Sex = F) И (Pclass = Low) И (Fare > 51.74))\n",
      "Proba=23 (Если (Sex = F) И (Pclass = High) И (Fare > 20.8))\n",
      "Proba=58 (Если (Sex = M) И (Cabin > 6.0) И (Age <= 3.5))\n",
      "Proba=94 (Если (Sex = M) И (Cabin <= 6.0) И (Age <= 14.0))\n"
     ]
    }
   ],
   "source": [
    "rules = get_rules_for_y_proba(model, df, 3)\n",
    "rules_str = format_rules_as_str(rules)\n",
    "rules_str = fix_comparison(rules_str)\n",
    "print(rules_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "041cccb8-fcd8-4150-be63-2f746060fbdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Если (Sex = M) И (Pclass = High) Тогда Важность_0=-0.921\n",
      "Если (Sex = M) И (Pclass = Low) Тогда Важность_0=-1.273\n",
      "Если (Sex = F) И (Pclass = Low) Тогда Важность_0=2.294\n",
      "Если (Sex = F) И (Pclass = High) Тогда Важность_0=1.495\n",
      "Если (Age > 8.5) И (Age <= 52.5) Тогда Важность_1=-0.172\n",
      "Если (Age > 8.5) И (Age > 52.5) Тогда Важность_1=-1.182\n",
      "Если (Age <= 8.5) И (SibSp <= 2.5) Тогда Важность_1=3.143\n",
      "Если (Age <= 8.5) И (SibSp > 2.5) Тогда Важность_1=1.459\n",
      "Если (Pclass = High) И (Sex = M) Тогда Важность_2=-0.466\n",
      "Если (Pclass = Low) И (Pclass <= 1.5) Тогда Важность_2=1.018\n",
      "Если (Pclass = Low) И (Pclass > 1.5) Тогда Важность_2=0.194\n",
      "Если (Pclass = High) И (Sex = F) Тогда Важность_2=-1.042\n"
     ]
    }
   ],
   "source": [
    "rules = get_rules_for_features(model, df, ['Sex', 'Age', 'Pclass'], 2)\n",
    "rules_str = format_rules_as_str(rules)\n",
    "rules_str = fix_comparison(rules_str)\n",
    "print(rules_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aabded4e-6659-48e0-adbe-b886600da668",
   "metadata": {},
   "source": [
    "Proba=12 (Если (Sex = M) И (Cabin > 6.0) И (Age > 3.5))\r\n",
    "Proba=41 (Если (Sex = M) И (Cabin <= 6.0) И (Age > 14.0))\r\n",
    "Proba=88 (Если (Sex = F) И (Pclass = Low) И (Fare <= 51.74))\r\n",
    "Proba=56 (Если (Sex = F) И (Pclass = High) И (Fare <= 20.8))\r\n",
    "Proba=97 (Если (Sex = F) И (Pclass = Low) И (Fare > 51.74))\r\n",
    "Proba=23 (Если (Sex = F) И (Pclass = High) И (Fare > 20.8))\r\n",
    "Proba=58 (Если (Sex = M) И (Cabin > 6.0) И (Age <= 3.5))\r\n",
    "Proba=94 (Если (Sex = M) И (Cabin <= 6.0) И (Age <= 14.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02e0c84d-a4e3-433f-82a8-660c014cd4a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "Proba = 96% Если ((Sex = F) И (Pclass = Low) И (Fare > 51.74));\n",
    "Proba = 93% Если ((Sex = M) И (Cabin <= 6.0) И (Age <= 14.0));\n",
    "Proba = 88% Если ((Sex = F) И (Pclass = Low) И (Fare <= 51.74));\n",
    "Proba = 58% Если ((Sex = M) И (Cabin > 6.0) И (Age <= 3.5));\n",
    "Proba = 56% Если ((Sex = F) И (Pclass = High) И (Fare <= 20.8));\n",
    "Proba = 12% Если ((Sex = M) И (Cabin > 6.0) И (Age > 3.5));\n",
    "Proba = 41% Если ((Sex = M) И (Cabin <= 6.0) И (Age > 14.0));\n",
    "Proba = 23% Если ((Sex = F) И (Pclass = High) И (Fare > 20.8));\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
